import os
import cv2
import numpy as np
import sys
from PIL import Image
from common.config import config_option

opencv_haarcascades_file = "D:\\directory\\opencv\\sources\\data\\haarcascades\\haarcascade_frontalface_default.xml"


def get_images_and_labels(path):
    image_paths = []
    ids = []
    detector = cv2.CascadeClassifier(opencv_haarcascades_file)

    for image_dir in os.listdir(path):
        image_file_path = path + os.sep + image_dir
        for image_file in os.listdir(image_file_path):
            image_paths.append(os.path.join(image_file_path, image_file))

    face_samples = []
    count = 0
    for image_path in image_paths:
        pil_img = Image.open(image_path).convert('L')
        img_numpy = np.array(pil_img, 'uint8')
        faces = detector.detectMultiScale(img_numpy)
        for (x, y, w, h) in faces:
            count += 1
            face_samples.append(img_numpy[y:y + h, x: x + w])
            ids.append(count)

    return face_samples, ids


def run_train_face_recognition():
    """运行人脸识别训练数据"""
    path = config_option['project_path'] + "/data/face_recognition/ORL_Faces"
    faces, ids = get_images_and_labels(path)
    recognizer = cv2.face.LBPHFaceRecognizer_create()
    recognizer.train(faces, np.array(ids))
    recognizer.write(config_option['project_path'] + "/output/face_recognition/train_face_recognition.yml")


if __name__ == '__main__':
    run_train_face_recognition()
